package zhipu

import (
	"chatgpt/ai/models"
	"chatgpt/helper"
	"chatgpt/helper/code"
	"chatgpt/logq"
	"chatgpt/plugin/httpsrv"
	"github.com/slclub/go-tips"
	"github.com/slclub/go-tips/json"
	"github.com/slclub/go-tips/stringbyte"
	"io"
	"net/http"
	"strconv"
	"strings"
)

var Search models.Search

func init() {
	Search.SystemWord = GLM_PROMPT_PRIFIX
}

func RemoteApi(w http.ResponseWriter, url_path string, m map[string]any, dest io.Writer) (code.CODENO, *helper.ResponseDesc) {
	cnf := helper.Config().Sub("Server")
	//addr_url := cnf.GetString("RmmoteAiUrl")

	res, resok := dest.(*helper.ResponseDesc)

	api_key := strings.ReplaceAll(helper.Config().Sub("zhipu").GetString("api_key"), "#", "%23")
	token, err := ConvKeyToToken(api_key)
	if err != nil {
		return code.ZHIPU_KEY_ERROR, nil
	}
	//param := fmt.Sprintf("?api_key=%v%v", strings.ReplaceAll(cnf.GetString("AuthToken"), "#", "%23"), pa1)
	domain := cnf.Sub("zhipu").GetString("RemoteAi")
	if domain == "" {
		domain = DOMAIN
	}
	url := domain + url_path
	url += tips.String(m["smodel"])
	url += "/"
	url += streamMethod(m["stream"])

	var ww http.ResponseWriter
	handleRead := httpsrv.HandleBody
	explainfunc := explainRemoteResponse
	ress := res
	if val, ok := m["stream"]; ok && (val == "true" || val == true) {
		handleRead = httpsrv.HandleStreamBody2
		explainfunc = explainStreamData
		ww = w
		//ress = nil
	}

	post := httpsrv.NewHttpPostStreamBoth(ww, url, m, ress).Head(func(requestH, responseH http.Header) {
		requestH.Set("Authorization", token)
		requestH.Set("Content-Type", "application/json; charset=utf-8")
	}).SetHandleRead(handleRead)
	post.ExplainDataFunc = explainfunc
	err = post.Handle()

	if err != nil {
		logq.Error("zhipu.ChatGLM.Remote err: ", err)
		res.MsgCode = code.ZHIPU_CHATGLM_ERROR
		return code.ZHIPU_CHATGLM_ERROR, res
	}

	if !resok {
		return code.SUCCESS, &helper.ResponseDesc{
			MsgCode: code.SUCCESS,
		}
	}
	//data, tokens := explainRemoteResponse(stringbyte.StringToBytes(res.Data.(string)))
	//res.Data = data
	//res.Tokens = tokens
	if res.MsgCode != code.SUCCESS {
		return res.MsgCode, res
	}
	//data := res.Data
	return res.MsgCode, res
}

func CompletionRemote(w http.ResponseWriter, url_path string, q string, m map[string]string) (code.CODENO, *helper.ResponseDesc) {
	res := helper.ResponseDesc{MsgCode: code.SUCCESS}
	return RemoteApi(w, url_path, remotePrefix(m), &res)
}

func CompletionRemoteString(w http.ResponseWriter, url_path string, q string, m map[string]string) (code.CODENO, *helper.ResponseDesc) {
	return RemoteApi(w, url_path, remotePrefix(m), nil)
}

func remotePrefix(m map[string]string) map[string]any {
	cnf := helper.Config().Sub("zhipu").Sub("completion")
	m_any := map[string]any{}
	for k, v := range m {
		m_any[k] = v
	}
	m_any["temperature"], _ = strconv.ParseFloat(cnf.GetString("temperature"), 64)
	m_any["top_p"], _ = strconv.ParseFloat(cnf.GetString("top_p"), 64)
	m_any["incremental"] = (cnf.GetBool("incremental"))

	if _, ok := m["zhishiku"]; !ok {
		m_any["zhishiku"] = cnf.GetString("zhishiku")
	}
	m_any["prompt"] = []map[string]any{{
		"role":    "user",
		"content": m["prompt"],
	}}
	return m_any
}

func streamMethod(me any) string {
	if tips.String(me) == "" {
		return "invoke"
	}
	return "sse-invoke"
}

func explainRemoteResponse(resp []byte) (rtn string, tokens int) {
	result := helper.ResponseDesc{}
	err := json.Unmarshal(resp, &result)
	if err != nil {
		return
	}
	data_map, ok := result.Data.(map[string]any)
	if !ok {
		return
	}
	choices, ok := data_map["choices"].([]any)
	if !ok || len(choices) < 1 {
		return
	}
	data, ok := choices[0].(map[string]any)
	if !ok {
		return
	}
	rtn = tips.String(data["content"])
	deststr := ""
	if errr := json.Unmarshal(stringbyte.StringToBytes(rtn), &deststr); errr == nil {
		rtn = deststr
	}
	//rtn = strings.ReplaceAll(rtn, "\"", "")

	usage, ok := data_map["usage"].(map[string]any)
	if !ok {
		return
	}
	tokens = tips.Int(usage["total_tokens"])
	return
}

func explainStreamData(b []byte) (string, int) {
	data := stringbyte.BytesToString(b)
	if tips.StrPos(data, "data:") != 0 {
		return "", 0
	}
	return string(b[5:]), len(b) - 5
}
